{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import pickle\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import pickle\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.utils import shuffle\n",
    "from torch.optim import optimizer\n",
    "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset\n",
    "from torch.nn import CrossEntropyLoss,BCEWithLogitsLoss\n",
    "from tqdm import tqdm_notebook, trange\n",
    "from pytorch_pretrained_bert import BertTokenizer, BertModel, BertForMaskedLM, BertForSequenceClassification\n",
    "from pytorch_pretrained_bert.optimization import BertAdam, WarmupLinearSchedule\n",
    "from sklearn.metrics import precision_recall_curve,classification_report\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('toutiao_cat_data.txt', 'r', encoding='utf8') as f:\n",
    "    data = f.read().splitlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "label = []\n",
    "title = []\n",
    "keyword = []\n",
    "\n",
    "for i in range(len(data)):\n",
    "    label.append(data[i].split('_!_')[1])\n",
    "    title.append(data[i].split('_!_')[3])\n",
    "    keyword.append(data[i].split('_!_')[4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>title</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>现在为什么都没人用千元机了？</td>\n",
       "      <td>109</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>农村小哥告诉你：5月份有哪些营养价值高的水果，不花钱随便吃</td>\n",
       "      <td>115</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>房地产翻倍涨价，实体经济却不好做，怎样看待这件事？</td>\n",
       "      <td>106</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>NBA2017年新秀在本赛季的实力怎么样？</td>\n",
       "      <td>103</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>为何移动支付在外国流行不起来，看完秒懂</td>\n",
       "      <td>109</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                           title label\n",
       "0                 现在为什么都没人用千元机了？   109\n",
       "1  农村小哥告诉你：5月份有哪些营养价值高的水果，不花钱随便吃   115\n",
       "2      房地产翻倍涨价，实体经济却不好做，怎样看待这件事？   106\n",
       "3          NBA2017年新秀在本赛季的实力怎么样？   103\n",
       "4            为何移动支付在外国流行不起来，看完秒懂   109"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "keyword = pd.DataFrame(keyword)\n",
    "title = pd.DataFrame(title)\n",
    "label = pd.DataFrame(label)\n",
    "trainData = pd.concat( [title , keyword], axis=1 )\n",
    "trainData = pd.concat( [trainData , label], axis=1 )\n",
    "trainData.columns = ['title', 'keyword', 'label']\n",
    "trainData = shuffle(trainData)\n",
    "\n",
    "trainData = trainData[:20000]\n",
    "trainData.index = range(len(trainData)) # row index從0開始\n",
    "trainData = trainData.drop(['keyword'], axis=1)\n",
    "trainData.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 對dataset label 做編碼\n",
    "le = LabelEncoder()\n",
    "trainData.label = le.fit_transform(trainData.label).reshape(-1,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>title</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>现在为什么都没人用千元机了？</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>农村小哥告诉你：5月份有哪些营养价值高的水果，不花钱随便吃</td>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>房地产翻倍涨价，实体经济却不好做，怎样看待这件事？</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>NBA2017年新秀在本赛季的实力怎么样？</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>为何移动支付在外国流行不起来，看完秒懂</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                           title  label\n",
       "0                 现在为什么都没人用千元机了？      8\n",
       "1  农村小哥告诉你：5月份有哪些营养价值高的水果，不花钱随便吃     13\n",
       "2      房地产翻倍涨价，实体经济却不好做，怎样看待这件事？      5\n",
       "3          NBA2017年新秀在本赛季的实力怎么样？      3\n",
       "4            为何移动支付在外国流行不起来，看完秒懂      8"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainData.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 随机选择20%作为测试集，剩余作为训练集\n",
    "x_train, x_test, y_train, y_test = train_test_split(trainData.title, trainData.label, test_size=0.2, random_state=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train = pd.DataFrame(x_train)\n",
    "x_test = pd.DataFrame(x_test)\n",
    "y_train = pd.DataFrame(y_train)\n",
    "y_test = pd.DataFrame(y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = pd.concat( [x_train , y_train], axis=1 )   # axis=1是X軸，axis=0是y軸\n",
    "test = pd.concat( [x_test , y_test], axis=1 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>title</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>17815</th>\n",
       "      <td>如何评价回国四子鹿晗、吴亦凡、黄子韬与张艺兴在选秀类综艺节目的表现？</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18370</th>\n",
       "      <td>古装剧里让人惋惜的5个角色，若能复活一个，你选谁？</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1379</th>\n",
       "      <td>驾考学车科目二倒车入库不会看角度，你不挂才怪</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14763</th>\n",
       "      <td>有没有首付15万以内，月供一千多的房源？</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7346</th>\n",
       "      <td>为什么说市面上很多的茅台都是假的？</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                    title  label\n",
       "17815  如何评价回国四子鹿晗、吴亦凡、黄子韬与张艺兴在选秀类综艺节目的表现？      2\n",
       "18370           古装剧里让人惋惜的5个角色，若能复活一个，你选谁？      2\n",
       "1379               驾考学车科目二倒车入库不会看角度，你不挂才怪      6\n",
       "14763                有没有首付15万以内，月供一千多的房源？      5\n",
       "7346                    为什么说市面上很多的茅台都是假的？      4"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 0 in train.label.values  查看0是否在label裡面\n",
    "train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>title</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>19134</th>\n",
       "      <td>你认为游戏中最孤独的角色是哪一个？</td>\n",
       "      <td>14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4981</th>\n",
       "      <td>高分必备！如何合理的分配试讲时间？</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16643</th>\n",
       "      <td>「《创文故事》纪录短片入围作品展播之⑭」书梦</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19117</th>\n",
       "      <td>为什么意大利人秃子这么多，中国人却很少？</td>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5306</th>\n",
       "      <td>散户应该怎么在股市上长期发展？</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        title  label\n",
       "19134       你认为游戏中最孤独的角色是哪一个？     14\n",
       "4981        高分必备！如何合理的分配试讲时间？      7\n",
       "16643  「《创文故事》纪录短片入围作品展播之⑭」书梦      1\n",
       "19117    为什么意大利人秃子这么多，中国人却很少？     11\n",
       "5306          散户应该怎么在股市上长期发展？      4"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The pre-trained model you are loading is an uncased model but you have set `do_lower_case` to False. We are setting `do_lower_case=True` for you but you may want to check this behavior.\n"
     ]
    }
   ],
   "source": [
    "# 分词工具\n",
    "bert_tokenizer = BertTokenizer.from_pretrained('bert-base-chinese', do_lower_case=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 封装类\n",
    "class DataPrecessForSingleSentence(object):\n",
    "    \"\"\"\n",
    "    对文本进行处理\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, bert_tokenizer, max_workers=10):\n",
    "        \"\"\"\n",
    "        bert_tokenizer :分词器\n",
    "        dataset        :包含列名为'text'与'label'的pandas dataframe\n",
    "        \"\"\"\n",
    "        self.bert_tokenizer = bert_tokenizer\n",
    "        # 创建多线程池\n",
    "        self.pool = ThreadPoolExecutor(max_workers=max_workers)\n",
    "        # 获取文本与标签\n",
    "\n",
    "    def get_input(self, dataset, max_seq_len=30):\n",
    "        \"\"\"\n",
    "        通过多线程（因为notebook中多进程使用存在一些问题）的方式对输入文本进行分词、ID化、截断、填充等流程得到最终的可用于模型输入的序列。\n",
    "        \n",
    "        入参:\n",
    "            dataset     : pandas的dataframe格式，包含两列，第一列为文本，第二列为标签。标签取值为{0,1}，其中0表示负样本，1代表正样本。\n",
    "            max_seq_len : 目标序列长度，该值需要预先对文本长度进行分别得到，可以设置为小于等于512（BERT的最长文本序列长度为512）的整数。\n",
    "        \n",
    "        出参:\n",
    "            seq         : 在入参seq的头尾分别拼接了'CLS'与'SEP'符号，如果长度仍小于max_seq_len，则使用0在尾部进行了填充。\n",
    "            seq_mask    : 只包含0、1且长度等于seq的序列，用于表征seq中的符号是否是有意义的，如果seq序列对应位上为填充符号，\n",
    "                          那么取值为1，否则为0。\n",
    "            seq_segment : shape等于seq，因为是单句，所以取值都为0。\n",
    "            labels      : 标签取值为{0,1}，其中0表示负样本，1代表正样本。\n",
    "        \n",
    "            \n",
    "        \"\"\"\n",
    "        sentences = dataset.iloc[:, 0].tolist()\n",
    "        labels = dataset.iloc[:, 1].tolist()\n",
    "        # 切词\n",
    "        tokens_seq = list(\n",
    "            self.pool.map(self.bert_tokenizer.tokenize, sentences))\n",
    "        # 获取定长序列及其mask\n",
    "        result = list(\n",
    "            self.pool.map(self.trunate_and_pad, tokens_seq,\n",
    "                          [max_seq_len] * len(tokens_seq)))\n",
    "        seqs = [i[0] for i in result]\n",
    "        seq_masks = [i[1] for i in result]\n",
    "        seq_segments = [i[2] for i in result]\n",
    "        return seqs, seq_masks, seq_segments, labels\n",
    "\n",
    "    def trunate_and_pad(self, seq, max_seq_len):\n",
    "        \"\"\"\n",
    "        1. 因为本类处理的是单句序列，按照BERT中的序列处理方式，需要在输入序列头尾分别拼接特殊字符'CLS'与'SEP'，\n",
    "           因此不包含两个特殊字符的序列长度应该小于等于max_seq_len-2，如果序列长度大于该值需要那么进行截断。\n",
    "        2. 对输入的序列 最终形成['CLS',seq,'SEP']的序列，该序列的长度如果小于max_seq_len，那么使用0进行填充。\n",
    "        \n",
    "        入参: \n",
    "            seq         : 输入序列，在本处其为单个句子。\n",
    "            max_seq_len : 拼接'CLS'与'SEP'这两个特殊字符后的序列长度\n",
    "        \n",
    "        出参:\n",
    "            seq         : 在入参seq的头尾分别拼接了'CLS'与'SEP'符号，如果长度仍小于max_seq_len，则使用0在尾部进行了填充。\n",
    "            seq_mask    : 只包含0、1且长度等于seq的序列，用于表征seq中的符号是否是有意义的，如果seq序列对应位上为填充符号，\n",
    "                          那么取值为1，否则为0。\n",
    "            seq_segment : shape等于seq，因为是单句，所以取值都为0。\n",
    "           \n",
    "        \"\"\"\n",
    "        # 对超长序列进行截断\n",
    "        if len(seq) > (max_seq_len - 2):\n",
    "            seq = seq[0:(max_seq_len - 2)]\n",
    "        # 分别在首尾拼接特殊符号\n",
    "        seq = ['[CLS]'] + seq + ['[SEP]']\n",
    "        # ID化\n",
    "        seq = self.bert_tokenizer.convert_tokens_to_ids(seq)\n",
    "        # 根据max_seq_len与seq的长度产生填充序列\n",
    "        padding = [0] * (max_seq_len - len(seq))\n",
    "        # 创建seq_mask\n",
    "        seq_mask = [1] * len(seq) + padding\n",
    "        # 创建seq_segment\n",
    "        seq_segment = [0] * len(seq) + padding\n",
    "        # 对seq拼接填充序列\n",
    "        seq += padding\n",
    "        assert len(seq) == max_seq_len\n",
    "        assert len(seq_mask) == max_seq_len\n",
    "        assert len(seq_segment) == max_seq_len\n",
    "        return seq, seq_mask, seq_segment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 类初始化\n",
    "processor = DataPrecessForSingleSentence(bert_tokenizer= bert_tokenizer)\n",
    "# 产生输入ju 数据\n",
    "seqs, seq_masks, seq_segments, labels = processor.get_input(\n",
    "    dataset=train, max_seq_len=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载预训练的bert模型\n",
    "model = BertForSequenceClassification.from_pretrained('bert-base-chinese', num_labels=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train转换为torch tensor\n",
    "t_seqs = torch.tensor(seqs, dtype=torch.long)\n",
    "t_seq_masks = torch.tensor(seq_masks, dtype = torch.long)\n",
    "t_seq_segments = torch.tensor(seq_segments, dtype = torch.long)\n",
    "t_labels = torch.tensor(labels, dtype = torch.long)\n",
    "\n",
    "train_data = TensorDataset(t_seqs, t_seq_masks, t_seq_segments, t_labels)\n",
    "train_sampler = RandomSampler(train_data)\n",
    "train_dataloder = DataLoader(dataset= train_data, sampler= train_sampler,batch_size = 256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test转换为torch tensor\n",
    "test_seqs, test_seq_masks, test_seq_segments, test_labels = processor.get_input(dataset=test, max_seq_len=50)\n",
    "\n",
    "test_seqs = torch.tensor(test_seqs, dtype=torch.long)\n",
    "test_seq_masks = torch.tensor(test_seq_masks, dtype = torch.long)\n",
    "test_seq_segments = torch.tensor(test_seq_segments, dtype = torch.long)\n",
    "test_labels = torch.tensor(test_labels, dtype = torch.long)\n",
    "test_data = TensorDataset(test_seqs, test_seq_masks, test_seq_segments, test_labels)\n",
    "test_dataloder = DataLoader(dataset= train_data, batch_size = 256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BertForSequenceClassification(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(21128, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): BertLayerNorm()\n",
       "      (dropout): Dropout(p=0.1)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (1): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (2): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (3): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (4): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (5): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (6): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (7): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (8): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (9): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (10): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "        (11): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): BertLayerNorm()\n",
       "              (dropout): Dropout(p=0.1)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): BertLayerNorm()\n",
       "            (dropout): Dropout(p=0.1)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pooler): BertPooler(\n",
       "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (activation): Tanh()\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1)\n",
       "  (classifier): Linear(in_features=768, out_features=15, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 将模型转换为trin mode\n",
    "model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# link : https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/pytorch_pretrained_bert/modeling.py\n",
    "class BertPooler(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super(BertPooler, self).__init__()\n",
    "        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n",
    "        self.activation = nn.Tanh()\n",
    "\n",
    "    def forward(self, hidden_states):\n",
    "        # We \"pool\" the model by simply taking the hidden state corresponding\n",
    "        # to the first token.\n",
    "        first_token_tensor = hidden_states[:, 0]\n",
    "        pooled_output = self.dense(first_token_tensor)\n",
    "        pooled_output = self.activation(pooled_output)\n",
    "        return pooled_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 待优化的参数\n",
    "param_optimizer = list(model.named_parameters())\n",
    "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n",
    "\n",
    "optimizer_grouped_parameters = [\n",
    "    {\n",
    "        'params':\n",
    "        [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],\n",
    "        'weight_decay':\n",
    "        0.01\n",
    "    },\n",
    "    {\n",
    "        'params':\n",
    "        [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],\n",
    "        'weight_decay':\n",
    "        0.0\n",
    "    }\n",
    "]\n",
    "\n",
    "optimizer = BertAdam(optimizer_grouped_parameters,\n",
    "                     lr=2e-05,\n",
    "                     warmup= 0.1 ,\n",
    "                     t_total= 2000)\n",
    "\n",
    "device = 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "Epoch:   0%|          | 0/10 [00:00<?, ?it/s]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7de813ba03054bd8bf2bcc155b04ab14",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Iteration', max=63), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 存储每一个batch的loss\n",
    "loss_collect = []\n",
    "for i in trange(10, desc='Epoch'):\n",
    "    for step, batch_data in enumerate(\n",
    "            tqdm_notebook(train_dataloder, desc='Iteration')):\n",
    "        batch_data = tuple(t.to(device) for t in batch_data)\n",
    "        batch_seqs, batch_seq_masks, batch_seq_segments, batch_labels = batch_data\n",
    "        # 对标签进行onehot编码\n",
    "        one_hot = torch.zeros(batch_labels.size(0), 28).long()\n",
    "        one_hot_batch_labels = one_hot.scatter_(\n",
    "            dim=1,\n",
    "            index=torch.unsqueeze(batch_labels, dim=1),\n",
    "            src=torch.ones(batch_labels.size(0), 28).long())\n",
    "\n",
    "        logits = model(\n",
    "            batch_seqs, batch_seq_masks, batch_seq_segments, labels=None)\n",
    "        logits = logits.softmax(dim=1)\n",
    "        loss_function = CrossEntropyLoss()\n",
    "        loss = loss_function(logits, batch_labels)\n",
    "        loss.backward()\n",
    "        loss_collect.append(loss.item())\n",
    "        print(\"\\r%f\" % loss, end='')\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,8))\n",
    "plt.plot(range(len(loss_collect)), loss_collect,'g.')\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model,open(\"fine_tuned_chinese_bert.bin\",\"wb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 用于存储预测标签与真实标签\n",
    "true_labels = []\n",
    "pred_labels = []\n",
    "model.eval()\n",
    "# 预测\n",
    "with torch.no_grad():\n",
    "    for batch_data in tqdm_notebook(test_dataloder, desc = 'TEST'):\n",
    "        batch_data = tuple(t.to(device) for t in batch_data)\n",
    "        batch_seqs, batch_seq_masks, batch_seq_segments, batch_labels = batch_data        \n",
    "        logits = model(\n",
    "            batch_seqs, batch_seq_masks, batch_seq_segments, labels=None)\n",
    "        logits = logits.softmax(dim=1).argmax(dim = 1)\n",
    "        pred_labels.append(logits.detach().numpy())\n",
    "        true_labels.append(batch_labels.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 查看各个类别的准召\n",
    "print(classification_report(np.concatenate(true_labels), np.concatenate(pred_labels)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
